#ifndef TEST_CNN_MPI_MNIST
#define TEST_CNN_MPI_MNIST

#include "hypercube/hypercube_cryto_agent.hpp"
#include "hypercube/hypercube_hmm_engine.hpp"
#include "mpi/mpi_util.hpp"
#include <helib/helib.h>
#include <torch/torch.h>

#define MAX(X, Y) (X) > (Y) ? (X) : (Y)

static constexpr int COMM_WORLD_ROOT = 0;
static constexpr int MPI_Bcast_ROOT = 0;

torch::Tensor init_conv_kernel(const torch::Tensor &a);

torch::Tensor init_fc1_kernel(const torch::Tensor &a, int conv_channels);

torch::Tensor init_fc2_kernel(const torch::Tensor &a);

torch::Tensor im2matrix(const torch::Tensor &x, const torch::Tensor &filter,
                        int stride);

template <typename Scheme> struct encrypted_model;
template <> struct encrypted_model<helib::CKKS> {
  std::vector<std::vector<helib::Ctxt>> conv;
  std::vector<helib::Ctxt> fc1;
  std::shared_ptr<helib::Ctxt> fc2;
};
template <> struct encrypted_model<helib::BGV> {
  std::vector<std::vector<std::vector<helib::Ctxt>>> conv;
  std::vector<std::vector<helib::Ctxt>> fc1;
  std::vector<helib::Ctxt> fc2;
};

template <typename Scheme> struct encrypted_input;
template <> struct encrypted_input<helib::CKKS> {
  std::vector<std::vector<helib::Ctxt>> data;
};
template <> struct encrypted_input<helib::BGV> {
  std::vector<std::vector<std::vector<std::vector<helib::Ctxt>>>> data;
};

template <typename Scheme> struct encrypted_output;
template <> struct encrypted_output<helib::CKKS> {
  std::shared_ptr<helib::Ctxt> data;
};
template <> struct encrypted_output<helib::BGV> {
  std::vector<helib::Ctxt> data;
};

struct data_server {
  std::vector<std::vector<std::vector<helib::Ctxt>>> conv;
  std::vector<std::vector<helib::Ctxt>> fc1;
  std::vector<helib::Ctxt> fc2;
};

template <typename Scheme> class mpi_h_mnist_client {
public:
  mpi_h_mnist_client() = default;
  virtual ~mpi_h_mnist_client() = default;

  // Generate encrypted model
  virtual void init_hcnn(encrypted_model<helib::BGV> &ctxt_model,
                         const torch::nn::Conv2d &conv1,
                         const torch::nn::Linear &fc1,
                         const torch::nn::Linear &fc2) const = 0;

  // Generate encrypted image
  virtual void init_input(encrypted_input<Scheme> &ctxt_input,
                          const torch::Tensor &filter,
                          const torch::Tensor &x) const = 0;

  virtual void recover_result(at::Tensor &out,
                              const std::vector<std::vector<std::string>> &in,
                              MPI_Comm comm, int root) const = 0;
};

template <typename Scheme> class mpi_h_mnist_server {
protected:
  struct metric {
  public:
    double hmm_encoding_time = 0.0;
    // * forword_test_time
    //   * hconv_time
    //   * hsq1_time
    //   * hfc1_time
    //   * hsq2_time
    //   * hfc2_time
    double forword_test_time = 0.0;
    double hconv_time = 0.0;
    double hsq1_time = 0.0;
    double hfc1_time = 0.0;
    double hsq2_time = 0.0;
    double hfc2_time = 0.0;

    double final_noise = 0.0;
    double init_noise = 0.0;
    double conv_noise = 0.0;
    double sq1_noise = 0.0;
    double fc1_noise = 0.0;
    double sq2_noise = 0.0;
    double fc2_noise = 0.0;

    double left_cap = 0.0;

  public:
    void reset() {
      // ! No need to reset `hmm_encoding_time`
      // ! One instance records once

      forword_test_time = 0.0;
      hconv_time = 0.0;
      hsq1_time = 0.0;
      hfc1_time = 0.0;
      hfc2_time = 0.0;
      hsq2_time = 0.0;
    }
  };

protected:
  mpi_h_mnist_server() = default;

public:
  MPI_Comm comm_reducfc1;
  MPI_Comm comm_reducfc2;
  MPI_Comm comm_Bcast;
  MPI_Comm comm_conv;

  virtual ~mpi_h_mnist_server() = default;

  virtual void forward_test(std::vector<std::vector<std::string>> &out,
                            std::vector<std::vector<std::string>> &buf_input,
                            const std::vector<std::vector<helib::Ctxt>> &hconv,
                            const std::vector<helib::Ctxt> &hfc1,
                            const std::vector<helib::Ctxt> &hfc2) const = 0;

  virtual void report_metrics(std::ostream &ss = std::cout) const = 0;

  virtual void bcast_hfc2_client(const std::vector<helib::Ctxt> &ctxt,
                                 MPI_Comm comm, std::vector<std::string> &buf,
                                 int root, int rank) const = 0;
  virtual void send_input_p2p(const encrypted_input<helib::BGV> &ctxt_in,
                              MPI_Comm comm,
                              std::vector<std::vector<std::string>> &buf,
                              int root) const = 0;

  virtual void
  send_honv_p2p(const std::vector<std::vector<std::vector<helib::Ctxt>>> &ctxt,
                MPI_Comm comm, std::vector<std::vector<std::string>> &buf,
                int root, int rank, int crt_num) const = 0;
  virtual void
  back_honv_p2p(const std::vector<std::vector<std::string>> &ctxt,
                MPI_Comm comm,
                std::vector<std::vector<std::vector<std::string>>> &buf,
                int root, int rank, int crt_num) const = 0;
  virtual void back_hfc1_p2p(const std::vector<std::string> &ctxt,
                             MPI_Comm comm,
                             std::vector<std::vector<std::string>> &buf,
                             int root, int rank, int crt_num) const = 0;
  virtual void send_hfc1_p2p(std::vector<std::vector<helib::Ctxt>> &ctxt,
                             MPI_Comm comm, std::vector<std::string> &buf,
                             int root, int rank, int crt_num) const = 0;
  virtual void
  bcast_hconv_server(const std::vector<std::vector<std::string>> &in,
                     MPI_Comm comm, std::vector<std::vector<std::string>> &buf,
                     int root, int rank, int crt_num) const = 0;
  virtual void
  back_hconv_server(const std::vector<std::vector<std::string>> &ctxt,
                    MPI_Comm comm,
                    std::vector<std::vector<std::vector<std::string>>> &buf,
                    int root, int rank, int crt_num) const = 0;
  virtual void
  back_input_server(const std::vector<std::vector<std::string>> &ctxt,
                    MPI_Comm comm,
                    std::vector<std::vector<std::vector<std::string>>> &buf,
                    int root, int rank, int crt_num) const = 0;
  virtual void bcast_hfc1_server(const std::vector<std::string> &in,
                                 MPI_Comm comm, std::vector<std::string> &buf,
                                 int root, int rank, int crt_num) const = 0;
};
class Net : public torch::nn::Module {
private:
  void ptxt_forward_test(const torch::Tensor &x);

public:
  Net(int root, MPI_Comm comm, const params<helib::BGV> &params);
  // ~Net() = default;
  ~Net() {
    delete client_;
    delete server_;
    // for (auto sk : sk_) {
    //   delete sk;
    // }
    // for (auto pk : pk_) {
    //   delete pk;
    // }
  }

  // const helib::PubKey &get_pk(int i) const { return *pk_[i]; }

  void CNN_init();

  torch::Tensor forward_train(torch::Tensor x);
  void init_in(const torch::Tensor &x, encrypted_input<helib::BGV> &ctxt_in);
  // torch::Tensor forward_test(const torch::Tensor &x);
  void forward_test(const torch::Tensor &x,
                    std::vector<std::vector<std::string>> &ctxt_out);
  torch::Tensor
  forward_result(const std::vector<std::vector<std::string>> &ctxt_out);
  double bcast_seckey(int root, int rank, MPI_Comm comm);

public:
  const mpi_h_mnist_client<helib::BGV> *client_ = nullptr;
  const mpi_h_mnist_server<helib::BGV> *server_ = nullptr;

  std::vector<std::vector<helib::Ctxt>> hconv;
  std::vector<helib::Ctxt> hfc1;
  std::vector<helib::Ctxt> hfc2;
  torch::nn::Conv2d conv1;
  torch::nn::Dropout2d conv1_drop;
  torch::nn::Linear fc1;
  torch::nn::Linear fc2;

  // std::vector<helib::SecKey *> sk_;
  // std::vector<helib::PubKey *> pk_;
};

// template <typename Scheme>
void test(Net &model, const torch::Device &device, const std::string &root,
          std::size_t test_num);

// template <typename Scheme>
void load_model(Net &model, const torch::Device &device,
                const std::string &pt_file);

#endif /* TEST_CNN_MNIST_HELPER */
